#!/usr/bin/python
import sys
import matplotlib.pyplot as pl
import numpy as np
from astro import aaomega
from astro.spec import plotlines,plotatmos

help =     """
right arrow, space:      Next spectrum
left arrow, backspace:   Prev spectrum
j:                       Enter spectrum index to jump to
y: Ly-alpha 1216
3: CIII     1908
4: CIV      1549
m: MgII     2798
h: H-beta   4862
a: H-alpha  6564 
"""

def plotconfig(spectra, info, outdir=None, show=True, index=0,
               comments=None, wmin=3700, wmax=8900, tweak_ymax=1.8):
    """ Plots fibres in a single configuration.

    Spectra and info are lists of spectra tuples, you can use this to
    plot more than one wavelength range for a given spectrum.

    can also give a subset of fibres to plot (fibres=), or save plots
    to a directory (outdir=)."""
    spectra = zip(*spectra)
    info = zip(*info)

    fig = pl.gcf()
    fig.clf()
    fig.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95,
                        wspace=0.2, hspace=0.2)
    for i in range(4):
        n = index + i
        cmap = 'brgmk'*(len(spectra[n]) / 5 + 1)
        sp = spectra[n][0]
        a = fig.add_subplot(4, 1, i + 1)
        aaomega.plotobject(spectra[n], info[n], ax=a, comments=None,
                           tweak_ymax=tweak_ymax, wmin=wmin, wmax=wmax)
        plotatmos(a)

class Junk(object):
    def __init__(self,spec,info):
        wmin = min(s[0].wa[0] for s in spec)
        wmax = max(s[0].wa[-1] for s in spec)
        self.wmin = wmin - 100
        self.wmax = wmax + 100
        self.i = 0
        self.spec = spec
        self.info = info
        self.n = len(spec[0])
        self.lines = []
        
    def update(self):
        self.lines = []
        plotconfig(self.spec, self.info, index=self.i,
                   wmin=self.wmin, wmax=self.wmax)
        pl.draw()
    def plotlines(self, zp1, ax):
        if self.lines:
            for l in self.lines:
                try:
                    l.remove()
                except ValueError:
                    # plot has been removed
                    pass
        self.lines = plotlines(zp1-1, ax)
        pl.draw()
    def on_keypress(self, event):
        if event.key in (' ','right'):
            if self.i >= self.n-4:
                print 'At last spectrum'
                return
            self.i += 4
            self.update()
        elif event.key in ('backspace','left'):
            if self.i == 0:
                print 'At first spectrum'
                return
            self.i -= 4
            if self.i < 0:  self.i = 0
            self.update()
        elif event.key == 'j':
            ind = raw_input('Jump to index (0-%i): ' % (self.n-1))
            self.i = int(ind)
            self.update()
        else:
            ax = event.inaxes
            if ax is None:  return
            if event.key == 'y':
                zp1 = event.xdata / 1215.6071
                print 'z=%.3f, Ly-alpha 1216' % (zp1-1)
            elif event.key == '4':
                zp1 = event.xdata / 1549.06
                print 'z=%.3f, CIV 1549' % (zp1-1)
            elif event.key == '3':
                zp1 = event.xdata / 1908.73
                print 'z=%.3f, CIII 1908' % (zp1-1)
            elif event.key == 'm':
                zp1 = event.xdata / 2798.75
                print 'z=%.3f, MgII 2798' % (zp1-1)
            elif event.key == 'h':
                zp1 = event.xdata / 4862.68
                print 'z=%.3f, H-beta 4862' % (zp1-1)
            elif event.key == 'a':
                zp1 = event.xdata / 6564.61
                print 'z=%.3f, H-alpha 6564' % (zp1-1)
            self.plotlines(zp1,ax)
    def connect(self,f):
        cids = list(f.canvas.callbacks.callbacks['key_press_event'])
        for cid in cids:
            f.canvas.callbacks.disconnect(cid)
        f.canvas.mpl_connect('key_press_event', self.on_keypress)

def main():
    spec = []
    info = []
    for name in sys.argv[1:]:
        print 'reading', name
        spec0,info0 = aaomega.readspec(name)
        spec.append(spec0)
        info.append(info0)

    junk = Junk(spec,info)
    fig = pl.figure()
    junk.connect(fig)
    junk.update()
    print help
    pl.show()

if __name__ == '__main__':
    main()
